import os
import json
import cv2
import numpy as np
import math
from PIL import Image, ImageDraw

class Rotate(object):

    def __init__(self, image: Image.Image, coordinate):
        self.image = image.convert('RGB')
        self.coordinate = coordinate
        self.xy = [tuple(self.coordinate[k]) for k in ['left_top', 'right_top', 'right_bottom', 'left_bottom']]
        self._mask = None
        self.image.putalpha(self.mask)

    @property
    def mask(self):
        if not self._mask:
            mask = Image.new('L', self.image.size, 0)
            draw = ImageDraw.Draw(mask, 'L')
            draw.polygon(self.xy, fill=255)
            self._mask = mask
        return self._mask

    def run(self):
        image = self.rotation_angle()
        box = image.getbbox()
        return image.crop(box)

    def rotation_angle(self):
        x1, y1 = self.xy[0]
        x2, y2 = self.xy[1]
        angle = self.angle([x1, y1, x2, y2], [0, 0, 10, 0]) * -1
        return self.image.rotate(angle, expand=True)

    def angle(self, v1, v2):
        dx1 = v1[2] - v1[0]
        dy1 = v1[3] - v1[1]
        dx2 = v2[2] - v2[0]
        dy2 = v2[3] - v2[1]
        angle1 = math.atan2(dy1, dx1)
        angle1 = int(angle1 * 180 / math.pi)
        angle2 = math.atan2(dy2, dx2)
        angle2 = int(angle2 * 180 / math.pi)
        if angle1 * angle2 >= 0:
            included_angle = abs(angle1 - angle2)
        else:
            included_angle = abs(angle1) + abs(angle2)
            if included_angle > 180:
                included_angle = 360 - included_angle
        return included_angle

with open('./M2021/Label.txt','r',encoding='utf8')as fp:
    s = [i[:-1].split('\t') for i in fp.readlines()]
    f1 = open('M2021_crop/rec_gt_train.txt', 'w', encoding='utf-8')
    f2 = open('M2021_crop/rec_gt_eval.txt', 'w', encoding='utf-8')
    for i in enumerate(s):
        path = i[1][0]
        anno = json.loads(i[1][1])
        filename = i[1][0][6:-4]
        image = Image.open(path)
        for j in range(len(anno)):
            label = anno[j]['transcription']
            roi = anno[j]['points']
            coordinate = {'left_top': anno[j]['points'][0], 'right_top': anno[j]['points'][1], 'right_bottom': anno[j]['points'][2], 'left_bottom': anno[j]['points'][3]}
            print(roi, label)
            rotate = Rotate(image, coordinate)
            # 把图片放到目录下
            crop_path = 'M2021_crop' + path[5:-4:] + '_' + str(j) + '.jpg'
            rotate.run().convert('RGB').save(crop_path)
            # label文件不写入图片目录
            crop_path = path[6:-4:] + '_' + str(j) + '.jpg'
            if i[0] % 5 != 0:
                f1.writelines(crop_path + '\t' + label + '\n')
            else:
                f2.writelines(crop_path + '\t' + label + '\n')
    f1.close()
    f2.close()